
#include "paging.h"
#include "memory/heap/kheap.h"
#include "status.h"
#include "config.h"

void paging_load_directory(uint32_t* directory); 

static uint32_t* current_directory = 0; 

struct paging_4Gbyte_chunk* paging_new_4Gbyte(uint8_t flags)
{
    uint32_t* directory = kzalloc(sizeof(uint32_t) * PAGING_TOTAL_ENTRIES_PER_TABLE); 
    int i;
    int j;  
    int offset = 0; 

    for (i = 0; i < PAGING_TOTAL_ENTRIES_PER_TABLE; i++) {
        uint32_t* entry = kzalloc(sizeof(uint32_t) * PAGING_TOTAL_ENTRIES_PER_TABLE); 
        
        for (j = 0; j < PAGING_TOTAL_ENTRIES_PER_TABLE; j++) {
            entry[j] = (offset + (j * PAGING_PAGE_SIZE)) | flags; 
        }

        offset += PAGING_TOTAL_ENTRIES_PER_TABLE * PAGING_PAGE_SIZE; 
        
        // TODO: check this how it exactly works
        directory[i] = (uint32_t)entry | flags | PAGING_IS_WRITEABLE; // address of entry that divides by 4096 (last 12 bits are 0), with flags, 
    }

    struct paging_4Gbyte_chunk* chunk = kzalloc(sizeof(struct paging_4Gbyte_chunk)); 
    chunk->directory_entry = directory; 
    return chunk; 
}

void paging_switch(struct paging_4Gbyte_chunk* directory)
{
    paging_load_directory(directory->directory_entry); 
    current_directory = directory->directory_entry; 
}

uint32_t* paging_4Gbyte_chunk_get_directory(struct paging_4Gbyte_chunk* chunk) 
{
    return chunk->directory_entry; 
}

int paging_is_aligned(void* address)
{
    return ((uint32_t)address % PAGING_PAGE_SIZE) == 0; 
}

int paging_get_indexes(void* virtual_address, uint32_t* directory_index_out, uint32_t* table_index_out)
{
    int res = 0; 
    if (!paging_is_aligned(virtual_address)) {   
        res = -TESTOS_STATUS_INVALID_ARG; 
        goto out; 
    }

    *directory_index_out = ((uint32_t)virtual_address / (PAGING_TOTAL_ENTRIES_PER_TABLE * PAGING_PAGE_SIZE)); 
    *table_index_out = ((uint32_t)virtual_address % (PAGING_TOTAL_ENTRIES_PER_TABLE * PAGING_PAGE_SIZE)) / PAGING_PAGE_SIZE; 
out:
    return res; 
}

int paging_set(uint32_t* directory, void* virtual_address, uint32_t value)
{ 
    uint32_t directory_index = 0; 
    uint32_t table_index = 0; 
    const int res = paging_get_indexes(virtual_address, &directory_index, &table_index); 
    if (res < 0) {
        goto out; 
    }
    
    const uint32_t entry = directory[directory_index]; 
    uint32_t* table = (uint32_t*)(entry & 0xFFFFF000); // extract address only, not flags, 31-11 bits are address bits
    table[table_index] = value; 

out:
    return res; 
}

void paging_free_4Gbyte(struct paging_4Gbyte_chunk* chunk)
{
    for (size_t i = 0; i < 1024; i++) {
        uint32_t entry = chunk->directory_entry[i]; 
        uint32_t* table = (uint32_t*)(entry & 0xFFFFF000); 
        kfree(table); 
    }

    kfree(chunk->directory_entry);
    kfree(chunk); 
}

const int paging_map(
    struct paging_4Gbyte_chunk* directory, 
    void* virt, 
    void* phys,
    const int flags)
{
    if (((unsigned int)virt % PAGING_PAGE_SIZE) || ((unsigned int)phys % PAGING_PAGE_SIZE)) {
        return TESTOS_STATUS_INVALID_ARG; 
    }

    return paging_set(directory->directory_entry, virt, (uint32_t)phys | (uint32_t)flags); 
}

const int paging_map_range(
    struct paging_4Gbyte_chunk* directory, 
    void* virt, 
    void* phys,
    const int count, 
    const int flags)
{
    int res = 0; 
    for (int i = 0; i < count; i++) {
        res = paging_map(directory, virt, phys, flags); 
        if (res < 0) {
            break; 
        }

        virt += PAGING_PAGE_SIZE; 
        phys += PAGING_PAGE_SIZE; 
    }

out:
    return res; 
}

const int paging_map_to(
    struct paging_4Gbyte_chunk* directory, 
    void* virt, 
    void* phys,
    void* phys_end,
    const int flags)
{
    int res = 0;
    if ((uint32_t)virt % PAGING_PAGE_SIZE) {
        res = -TESTOS_STATUS_INVALID_ARG; 
        goto out; 
    } 

    if ((uint32_t)phys % PAGING_PAGE_SIZE) {
        res = -TESTOS_STATUS_INVALID_ARG; 
        goto out; 
    } 

    if ((uint32_t)phys_end % PAGING_PAGE_SIZE) {
        res = -TESTOS_STATUS_INVALID_ARG; 
        goto out; 
    } 

    if ((uint32_t)phys_end < (uint32_t)phys)
    {
        res = -TESTOS_STATUS_INVALID_ARG; 
        goto out; 
    }

    const uint32_t total_bytes = (uint32_t)phys_end - (uint32_t)phys; 
    const int total_pages = total_bytes / PAGING_PAGE_SIZE; 
    res = paging_map_range(directory, virt, phys, total_pages, flags); 

out:
    return res; 
}

void* paging_align_address(void* ptr)
{
    uint32_t iptr = (uint32_t)ptr; 
    uint32_t mod = iptr % PAGING_PAGE_SIZE; 
    return (void*)(iptr + PAGING_PAGE_SIZE - mod); 
}

const uint32_t paging_get(uint32_t* directory, void* virt)
{
    uint32_t directory_index = 0; 
    uint32_t table_index = 0; 
    // int res = paging_get_indexes(virt, &directory_index, &table_index); 
    paging_get_indexes(virt, &directory_index, &table_index); 
    uint32_t entry = directory[directory_index]; 
    uint32_t* table = (uint32_t*)(entry & 0xFFFFF000); 
    return table[table_index]; 
}